import torch
import torch_npu

# 检查是否有可用的 NPU 设备
if torch_npu.npu.is_available():
    # 获取默认 NPU 设备
    device = torch_npu.npu.current_device()

    # 获取 NPU 专用内存信息
    total_memory = torch_npu.npu.get_device_properties(device).total_memory
    allocated_memory = torch_npu.npu.memory_allocated(device)
    cached_memory = torch_npu.npu.memory_reserved(device)
    free_memory = total_memory - allocated_memory

    print(f"Total NPU memory: {total_memory / 1024 ** 2:.2f} MB")
    print(f"Allocated NPU memory: {allocated_memory / 1024 ** 2:.2f} MB")
    print(f"Cached NPU memory: {cached_memory / 1024 ** 2:.2f} MB")
    print(f"Free NPU memory: {free_memory / 1024 ** 2:.2f} MB")
else:
    print("No NPU devices available.")